{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9bfd0840",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41ffaa67",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3633862/484069769.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/mtm/nan/3/both_bilstm.betterthanlast.b21_s15_p20_epoch19.ckpt\", map_location=\"cpu\")\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 1) 载入 checkpoint（不要从不信任来源直接执行）\n",
    "ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/mtm/nan/3/both_bilstm.betterthanlast.b21_s15_p20_epoch19.ckpt\", map_location=\"cpu\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2a1f38ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'collections.OrderedDict'>\n",
      "keys: ['cls_tok', 'rpe', 'ape', 'embedding.weight', 'chn_emb.weight', 'inp_layer.temporal.wq.weight', 'inp_layer.temporal.wq.bias', 'inp_layer.temporal.wk.weight', 'inp_layer.temporal.wk.bias', 'inp_layer.temporal.wv.weight', 'inp_layer.temporal.wv.bias', 'inp_layer.temporal.layer_norm.weight', 'inp_layer.temporal.layer_norm.bias', 'inp_layer.temporal.layer_scale.gamma', 'inp_layer.mixer.wq.weight', 'inp_layer.mixer.wq.bias', 'inp_layer.mixer.wk.weight', 'inp_layer.mixer.wk.bias', 'inp_layer.mixer.wv.weight', 'inp_layer.mixer.wv.bias']\n",
      "cls_tok (5, 128)\n",
      "rpe (640, 64)\n",
      "ape (640, 128)\n",
      "embedding.weight (16, 4)\n",
      "chn_emb.weight (5, 128)\n",
      "inp_layer.temporal.wq.weight (128, 128)\n",
      "inp_layer.temporal.wq.bias (128,)\n",
      "inp_layer.temporal.wk.weight (128, 128)\n",
      "inp_layer.temporal.wk.bias (128,)\n",
      "inp_layer.temporal.wv.weight (128, 128)\n",
      "inp_layer.temporal.wv.bias (128,)\n",
      "inp_layer.temporal.layer_norm.weight (128,)\n",
      "inp_layer.temporal.layer_norm.bias (128,)\n",
      "inp_layer.temporal.layer_scale.gamma (128,)\n",
      "inp_layer.mixer.wq.weight (128, 128)\n",
      "inp_layer.mixer.wq.bias (128,)\n",
      "inp_layer.mixer.wk.weight (128, 128)\n",
      "inp_layer.mixer.wk.bias (128,)\n",
      "inp_layer.mixer.wv.weight (128, 128)\n",
      "inp_layer.mixer.wv.bias (128,)\n",
      "inp_layer.mixer.layer_norm.weight (128,)\n",
      "inp_layer.mixer.layer_norm.bias (128,)\n",
      "inp_layer.mixer.layer_scale.gamma (128,)\n",
      "inp_layer.mlp3.net.0.weight (512, 128)\n",
      "inp_layer.mlp3.net.0.bias (512,)\n",
      "inp_layer.mlp3.net.2.weight (128, 512)\n",
      "inp_layer.mlp3.net.2.bias (128,)\n",
      "inp_layer.mlp3.net.3.gamma (128,)\n",
      "inp_layer.mlp3.norm.weight (128,)\n",
      "inp_layer.mlp3.norm.bias (128,)\n",
      "mixers.0.temporal.wq.weight (128, 128)\n",
      "mixers.0.temporal.wq.bias (128,)\n",
      "mixers.0.temporal.wk.weight (128, 128)\n",
      "mixers.0.temporal.wk.bias (128,)\n",
      "mixers.0.temporal.wv.weight (128, 128)\n",
      "mixers.0.temporal.wv.bias (128,)\n",
      "mixers.0.temporal.layer_norm.weight (128,)\n",
      "mixers.0.temporal.layer_norm.bias (128,)\n",
      "mixers.0.temporal.layer_scale.gamma (128,)\n",
      "mixers.0.mixer.wq.weight (128, 128)\n",
      "mixers.0.mixer.wq.bias (128,)\n",
      "mixers.0.mixer.wk.weight (128, 128)\n",
      "mixers.0.mixer.wk.bias (128,)\n",
      "mixers.0.mixer.wv.weight (128, 128)\n",
      "mixers.0.mixer.wv.bias (128,)\n",
      "mixers.0.mixer.layer_norm.weight (128,)\n",
      "mixers.0.mixer.layer_norm.bias (128,)\n",
      "mixers.0.mixer.layer_scale.gamma (128,)\n",
      "mixers.0.mlp3.net.0.weight (512, 128)\n",
      "mixers.0.mlp3.net.0.bias (512,)\n",
      "mixers.0.mlp3.net.2.weight (128, 512)\n",
      "mixers.0.mlp3.net.2.bias (128,)\n",
      "mixers.0.mlp3.net.3.gamma (128,)\n",
      "mixers.0.mlp3.norm.weight (128,)\n",
      "mixers.0.mlp3.norm.bias (128,)\n",
      "mixers.1.temporal.wq.weight (128, 128)\n",
      "mixers.1.temporal.wq.bias (128,)\n",
      "mixers.1.temporal.wk.weight (128, 128)\n",
      "mixers.1.temporal.wk.bias (128,)\n",
      "mixers.1.temporal.wv.weight (128, 128)\n",
      "mixers.1.temporal.wv.bias (128,)\n",
      "mixers.1.temporal.layer_norm.weight (128,)\n",
      "mixers.1.temporal.layer_norm.bias (128,)\n",
      "mixers.1.temporal.layer_scale.gamma (128,)\n",
      "mixers.1.mixer.wq.weight (128, 128)\n",
      "mixers.1.mixer.wq.bias (128,)\n",
      "mixers.1.mixer.wk.weight (128, 128)\n",
      "mixers.1.mixer.wk.bias (128,)\n",
      "mixers.1.mixer.wv.weight (128, 128)\n",
      "mixers.1.mixer.wv.bias (128,)\n",
      "mixers.1.mixer.layer_norm.weight (128,)\n",
      "mixers.1.mixer.layer_norm.bias (128,)\n",
      "mixers.1.mixer.layer_scale.gamma (128,)\n",
      "mixers.1.mlp3.net.0.weight (512, 128)\n",
      "mixers.1.mlp3.net.0.bias (512,)\n",
      "mixers.1.mlp3.net.2.weight (128, 512)\n",
      "mixers.1.mlp3.net.2.bias (128,)\n",
      "mixers.1.mlp3.net.3.gamma (128,)\n",
      "mixers.1.mlp3.norm.weight (128,)\n",
      "mixers.1.mlp3.norm.bias (128,)\n",
      "mixers.2.temporal.wq.weight (128, 128)\n",
      "mixers.2.temporal.wq.bias (128,)\n",
      "mixers.2.temporal.wk.weight (128, 128)\n",
      "mixers.2.temporal.wk.bias (128,)\n",
      "mixers.2.temporal.wv.weight (128, 128)\n",
      "mixers.2.temporal.wv.bias (128,)\n",
      "mixers.2.temporal.layer_norm.weight (128,)\n",
      "mixers.2.temporal.layer_norm.bias (128,)\n",
      "mixers.2.temporal.layer_scale.gamma (128,)\n",
      "mixers.2.mixer.wq.weight (128, 128)\n",
      "mixers.2.mixer.wq.bias (128,)\n",
      "mixers.2.mixer.wk.weight (128, 128)\n",
      "mixers.2.mixer.wk.bias (128,)\n",
      "mixers.2.mixer.wv.weight (128, 128)\n",
      "mixers.2.mixer.wv.bias (128,)\n",
      "mixers.2.mixer.layer_norm.weight (128,)\n",
      "mixers.2.mixer.layer_norm.bias (128,)\n",
      "mixers.2.mixer.layer_scale.gamma (128,)\n",
      "mixers.2.mlp3.net.0.weight (512, 128)\n",
      "mixers.2.mlp3.net.0.bias (512,)\n",
      "mixers.2.mlp3.net.2.weight (128, 512)\n",
      "mixers.2.mlp3.net.2.bias (128,)\n",
      "mixers.2.mlp3.net.3.gamma (128,)\n",
      "mixers.2.mlp3.norm.weight (128,)\n",
      "mixers.2.mlp3.norm.bias (128,)\n",
      "mixers.3.temporal.wq.weight (128, 128)\n",
      "mixers.3.temporal.wq.bias (128,)\n",
      "mixers.3.temporal.wk.weight (128, 128)\n",
      "mixers.3.temporal.wk.bias (128,)\n",
      "mixers.3.temporal.wv.weight (128, 128)\n",
      "mixers.3.temporal.wv.bias (128,)\n",
      "mixers.3.temporal.layer_norm.weight (128,)\n",
      "mixers.3.temporal.layer_norm.bias (128,)\n",
      "mixers.3.temporal.layer_scale.gamma (128,)\n",
      "mixers.3.mixer.wq.weight (128, 128)\n",
      "mixers.3.mixer.wq.bias (128,)\n",
      "mixers.3.mixer.wk.weight (128, 128)\n",
      "mixers.3.mixer.wk.bias (128,)\n",
      "mixers.3.mixer.wv.weight (128, 128)\n",
      "mixers.3.mixer.wv.bias (128,)\n",
      "mixers.3.mixer.layer_norm.weight (128,)\n",
      "mixers.3.mixer.layer_norm.bias (128,)\n",
      "mixers.3.mixer.layer_scale.gamma (128,)\n",
      "mixers.3.mlp3.net.0.weight (512, 128)\n",
      "mixers.3.mlp3.net.0.bias (512,)\n",
      "mixers.3.mlp3.net.2.weight (128, 512)\n",
      "mixers.3.mlp3.net.2.bias (128,)\n",
      "mixers.3.mlp3.net.3.gamma (128,)\n",
      "mixers.3.mlp3.norm.weight (128,)\n",
      "mixers.3.mlp3.norm.bias (128,)\n",
      "samplers.0.down.lin.weight (128, 256)\n",
      "samplers.0.down.lin.bias (128,)\n",
      "samplers.1.down.lin.weight (128, 256)\n",
      "samplers.1.down.lin.bias (128,)\n",
      "samplers.2.down.lin.weight (128, 256)\n",
      "samplers.2.down.lin.bias (128,)\n",
      "samplers.3.down.lin.weight (128, 256)\n",
      "samplers.3.down.lin.bias (128,)\n",
      "cls_head.net.0.weight (512, 129)\n",
      "cls_head.net.0.bias (512,)\n",
      "cls_head.net.3.weight (2, 512)\n",
      "cls_head.net.3.bias (2,)\n",
      "total params: 1483074\n"
     ]
    }
   ],
   "source": [
    "print(type(ckpt))\n",
    "\n",
    "# 2) 常见情况：Lightning 保存的 checkpoint 包含 'state_dict'\n",
    "if isinstance(ckpt, dict):\n",
    "    print(\"keys:\", list(ckpt.keys())[:20])\n",
    "    if \"state_dict\" in ckpt:\n",
    "        sd = ckpt[\"state_dict\"]\n",
    "    elif \"model_state_dict\" in ckpt:\n",
    "        sd = ckpt[\"model_state_dict\"]\n",
    "    else:\n",
    "        # 有些直接保存了 state_dict 本身\n",
    "        sd = ckpt\n",
    "\n",
    "    # 列出参数名和形状\n",
    "    for k, v in sd.items():\n",
    "        try:\n",
    "            print(k, tuple(v.shape))\n",
    "        except Exception:\n",
    "            print(k, type(v))\n",
    "\n",
    "    # 统计参数量\n",
    "    total = sum(v.numel() for v in sd.values() if hasattr(v, \"numel\"))\n",
    "    print(\"total params:\", total)\n",
    "else:\n",
    "    print(\"checkpoint is not a dict, maybe a model object:\", ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "28ffef3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3633862/3389481113.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/mtm/tmp/both_bilstm.betterthanlast.b21_s15_p20_epoch18.ckpt\", map_location=\"cpu\")\n"
     ]
    }
   ],
   "source": [
    "# 1) 载入 checkpoint（不要从不信任来源直接执行）\n",
    "ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/mtm/tmp/both_bilstm.betterthanlast.b21_s15_p20_epoch18.ckpt\", map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f68d2b51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'collections.OrderedDict'>\n",
      "keys: ['cls_tok', 'rpe', 'ape', 'embedding.weight', 'chn_emb.weight', 'inp_layer.temporal.wq.weight', 'inp_layer.temporal.wq.bias', 'inp_layer.temporal.wk.weight', 'inp_layer.temporal.wk.bias', 'inp_layer.temporal.wv.weight', 'inp_layer.temporal.wv.bias', 'inp_layer.temporal.layer_norm.weight', 'inp_layer.temporal.layer_norm.bias', 'inp_layer.temporal.layer_scale.gamma', 'inp_layer.mixer.wq.weight', 'inp_layer.mixer.wq.bias', 'inp_layer.mixer.wk.weight', 'inp_layer.mixer.wk.bias', 'inp_layer.mixer.wv.weight', 'inp_layer.mixer.wv.bias']\n",
      "cls_tok (9, 128)\n",
      "rpe (640, 64)\n",
      "ape (640, 128)\n",
      "embedding.weight (16, 8)\n",
      "chn_emb.weight (9, 128)\n",
      "inp_layer.temporal.wq.weight (128, 128)\n",
      "inp_layer.temporal.wq.bias (128,)\n",
      "inp_layer.temporal.wk.weight (128, 128)\n",
      "inp_layer.temporal.wk.bias (128,)\n",
      "inp_layer.temporal.wv.weight (128, 128)\n",
      "inp_layer.temporal.wv.bias (128,)\n",
      "inp_layer.temporal.layer_norm.weight (128,)\n",
      "inp_layer.temporal.layer_norm.bias (128,)\n",
      "inp_layer.temporal.layer_scale.gamma (128,)\n",
      "inp_layer.mixer.wq.weight (128, 128)\n",
      "inp_layer.mixer.wq.bias (128,)\n",
      "inp_layer.mixer.wk.weight (128, 128)\n",
      "inp_layer.mixer.wk.bias (128,)\n",
      "inp_layer.mixer.wv.weight (128, 128)\n",
      "inp_layer.mixer.wv.bias (128,)\n",
      "inp_layer.mixer.layer_norm.weight (128,)\n",
      "inp_layer.mixer.layer_norm.bias (128,)\n",
      "inp_layer.mixer.layer_scale.gamma (128,)\n",
      "inp_layer.channel.wq.weight (128, 128)\n",
      "inp_layer.channel.wq.bias (128,)\n",
      "inp_layer.channel.wk.weight (128, 128)\n",
      "inp_layer.channel.wk.bias (128,)\n",
      "inp_layer.channel.wv.weight (128, 128)\n",
      "inp_layer.channel.wv.bias (128,)\n",
      "inp_layer.channel.layer_norm.weight (128,)\n",
      "inp_layer.channel.layer_norm.bias (128,)\n",
      "inp_layer.channel.layer_scale.gamma (128,)\n",
      "inp_layer.mlp3.net.0.weight (512, 128)\n",
      "inp_layer.mlp3.net.0.bias (512,)\n",
      "inp_layer.mlp3.net.2.weight (128, 512)\n",
      "inp_layer.mlp3.net.2.bias (128,)\n",
      "inp_layer.mlp3.net.3.gamma (128,)\n",
      "inp_layer.mlp3.norm.weight (128,)\n",
      "inp_layer.mlp3.norm.bias (128,)\n",
      "mixers.0.temporal.wq.weight (128, 128)\n",
      "mixers.0.temporal.wq.bias (128,)\n",
      "mixers.0.temporal.wk.weight (128, 128)\n",
      "mixers.0.temporal.wk.bias (128,)\n",
      "mixers.0.temporal.wv.weight (128, 128)\n",
      "mixers.0.temporal.wv.bias (128,)\n",
      "mixers.0.temporal.layer_norm.weight (128,)\n",
      "mixers.0.temporal.layer_norm.bias (128,)\n",
      "mixers.0.temporal.layer_scale.gamma (128,)\n",
      "mixers.0.mixer.wq.weight (128, 128)\n",
      "mixers.0.mixer.wq.bias (128,)\n",
      "mixers.0.mixer.wk.weight (128, 128)\n",
      "mixers.0.mixer.wk.bias (128,)\n",
      "mixers.0.mixer.wv.weight (128, 128)\n",
      "mixers.0.mixer.wv.bias (128,)\n",
      "mixers.0.mixer.layer_norm.weight (128,)\n",
      "mixers.0.mixer.layer_norm.bias (128,)\n",
      "mixers.0.mixer.layer_scale.gamma (128,)\n",
      "mixers.0.channel.wq.weight (128, 128)\n",
      "mixers.0.channel.wq.bias (128,)\n",
      "mixers.0.channel.wk.weight (128, 128)\n",
      "mixers.0.channel.wk.bias (128,)\n",
      "mixers.0.channel.wv.weight (128, 128)\n",
      "mixers.0.channel.wv.bias (128,)\n",
      "mixers.0.channel.layer_norm.weight (128,)\n",
      "mixers.0.channel.layer_norm.bias (128,)\n",
      "mixers.0.channel.layer_scale.gamma (128,)\n",
      "mixers.0.mlp3.net.0.weight (512, 128)\n",
      "mixers.0.mlp3.net.0.bias (512,)\n",
      "mixers.0.mlp3.net.2.weight (128, 512)\n",
      "mixers.0.mlp3.net.2.bias (128,)\n",
      "mixers.0.mlp3.net.3.gamma (128,)\n",
      "mixers.0.mlp3.norm.weight (128,)\n",
      "mixers.0.mlp3.norm.bias (128,)\n",
      "mixers.1.temporal.wq.weight (128, 128)\n",
      "mixers.1.temporal.wq.bias (128,)\n",
      "mixers.1.temporal.wk.weight (128, 128)\n",
      "mixers.1.temporal.wk.bias (128,)\n",
      "mixers.1.temporal.wv.weight (128, 128)\n",
      "mixers.1.temporal.wv.bias (128,)\n",
      "mixers.1.temporal.layer_norm.weight (128,)\n",
      "mixers.1.temporal.layer_norm.bias (128,)\n",
      "mixers.1.temporal.layer_scale.gamma (128,)\n",
      "mixers.1.mixer.wq.weight (128, 128)\n",
      "mixers.1.mixer.wq.bias (128,)\n",
      "mixers.1.mixer.wk.weight (128, 128)\n",
      "mixers.1.mixer.wk.bias (128,)\n",
      "mixers.1.mixer.wv.weight (128, 128)\n",
      "mixers.1.mixer.wv.bias (128,)\n",
      "mixers.1.mixer.layer_norm.weight (128,)\n",
      "mixers.1.mixer.layer_norm.bias (128,)\n",
      "mixers.1.mixer.layer_scale.gamma (128,)\n",
      "mixers.1.channel.wq.weight (128, 128)\n",
      "mixers.1.channel.wq.bias (128,)\n",
      "mixers.1.channel.wk.weight (128, 128)\n",
      "mixers.1.channel.wk.bias (128,)\n",
      "mixers.1.channel.wv.weight (128, 128)\n",
      "mixers.1.channel.wv.bias (128,)\n",
      "mixers.1.channel.layer_norm.weight (128,)\n",
      "mixers.1.channel.layer_norm.bias (128,)\n",
      "mixers.1.channel.layer_scale.gamma (128,)\n",
      "mixers.1.mlp3.net.0.weight (512, 128)\n",
      "mixers.1.mlp3.net.0.bias (512,)\n",
      "mixers.1.mlp3.net.2.weight (128, 512)\n",
      "mixers.1.mlp3.net.2.bias (128,)\n",
      "mixers.1.mlp3.net.3.gamma (128,)\n",
      "mixers.1.mlp3.norm.weight (128,)\n",
      "mixers.1.mlp3.norm.bias (128,)\n",
      "mixers.2.temporal.wq.weight (128, 128)\n",
      "mixers.2.temporal.wq.bias (128,)\n",
      "mixers.2.temporal.wk.weight (128, 128)\n",
      "mixers.2.temporal.wk.bias (128,)\n",
      "mixers.2.temporal.wv.weight (128, 128)\n",
      "mixers.2.temporal.wv.bias (128,)\n",
      "mixers.2.temporal.layer_norm.weight (128,)\n",
      "mixers.2.temporal.layer_norm.bias (128,)\n",
      "mixers.2.temporal.layer_scale.gamma (128,)\n",
      "mixers.2.mixer.wq.weight (128, 128)\n",
      "mixers.2.mixer.wq.bias (128,)\n",
      "mixers.2.mixer.wk.weight (128, 128)\n",
      "mixers.2.mixer.wk.bias (128,)\n",
      "mixers.2.mixer.wv.weight (128, 128)\n",
      "mixers.2.mixer.wv.bias (128,)\n",
      "mixers.2.mixer.layer_norm.weight (128,)\n",
      "mixers.2.mixer.layer_norm.bias (128,)\n",
      "mixers.2.mixer.layer_scale.gamma (128,)\n",
      "mixers.2.channel.wq.weight (128, 128)\n",
      "mixers.2.channel.wq.bias (128,)\n",
      "mixers.2.channel.wk.weight (128, 128)\n",
      "mixers.2.channel.wk.bias (128,)\n",
      "mixers.2.channel.wv.weight (128, 128)\n",
      "mixers.2.channel.wv.bias (128,)\n",
      "mixers.2.channel.layer_norm.weight (128,)\n",
      "mixers.2.channel.layer_norm.bias (128,)\n",
      "mixers.2.channel.layer_scale.gamma (128,)\n",
      "mixers.2.mlp3.net.0.weight (512, 128)\n",
      "mixers.2.mlp3.net.0.bias (512,)\n",
      "mixers.2.mlp3.net.2.weight (128, 512)\n",
      "mixers.2.mlp3.net.2.bias (128,)\n",
      "mixers.2.mlp3.net.3.gamma (128,)\n",
      "mixers.2.mlp3.norm.weight (128,)\n",
      "mixers.2.mlp3.norm.bias (128,)\n",
      "mixers.3.temporal.wq.weight (128, 128)\n",
      "mixers.3.temporal.wq.bias (128,)\n",
      "mixers.3.temporal.wk.weight (128, 128)\n",
      "mixers.3.temporal.wk.bias (128,)\n",
      "mixers.3.temporal.wv.weight (128, 128)\n",
      "mixers.3.temporal.wv.bias (128,)\n",
      "mixers.3.temporal.layer_norm.weight (128,)\n",
      "mixers.3.temporal.layer_norm.bias (128,)\n",
      "mixers.3.temporal.layer_scale.gamma (128,)\n",
      "mixers.3.mixer.wq.weight (128, 128)\n",
      "mixers.3.mixer.wq.bias (128,)\n",
      "mixers.3.mixer.wk.weight (128, 128)\n",
      "mixers.3.mixer.wk.bias (128,)\n",
      "mixers.3.mixer.wv.weight (128, 128)\n",
      "mixers.3.mixer.wv.bias (128,)\n",
      "mixers.3.mixer.layer_norm.weight (128,)\n",
      "mixers.3.mixer.layer_norm.bias (128,)\n",
      "mixers.3.mixer.layer_scale.gamma (128,)\n",
      "mixers.3.channel.wq.weight (128, 128)\n",
      "mixers.3.channel.wq.bias (128,)\n",
      "mixers.3.channel.wk.weight (128, 128)\n",
      "mixers.3.channel.wk.bias (128,)\n",
      "mixers.3.channel.wv.weight (128, 128)\n",
      "mixers.3.channel.wv.bias (128,)\n",
      "mixers.3.channel.layer_norm.weight (128,)\n",
      "mixers.3.channel.layer_norm.bias (128,)\n",
      "mixers.3.channel.layer_scale.gamma (128,)\n",
      "mixers.3.mlp3.net.0.weight (512, 128)\n",
      "mixers.3.mlp3.net.0.bias (512,)\n",
      "mixers.3.mlp3.net.2.weight (128, 512)\n",
      "mixers.3.mlp3.net.2.bias (128,)\n",
      "mixers.3.mlp3.net.3.gamma (128,)\n",
      "mixers.3.mlp3.norm.weight (128,)\n",
      "mixers.3.mlp3.norm.bias (128,)\n",
      "samplers.0.down.lin.weight (128, 256)\n",
      "samplers.0.down.lin.bias (128,)\n",
      "samplers.1.down.lin.weight (128, 256)\n",
      "samplers.1.down.lin.bias (128,)\n",
      "samplers.2.down.lin.weight (128, 256)\n",
      "samplers.2.down.lin.bias (128,)\n",
      "samplers.3.down.lin.weight (128, 256)\n",
      "samplers.3.down.lin.bias (128,)\n",
      "cls_head.net.0.weight (512, 129)\n",
      "cls_head.net.0.bias (512,)\n",
      "cls_head.net.3.weight (2, 512)\n",
      "cls_head.net.3.bias (2,)\n",
      "total params: 1733762\n"
     ]
    }
   ],
   "source": [
    "print(type(ckpt))\n",
    "\n",
    "# 2) 常见情况：Lightning 保存的 checkpoint 包含 'state_dict'\n",
    "if isinstance(ckpt, dict):\n",
    "    print(\"keys:\", list(ckpt.keys())[:20])\n",
    "    if \"state_dict\" in ckpt:\n",
    "        sd = ckpt[\"state_dict\"]\n",
    "    elif \"model_state_dict\" in ckpt:\n",
    "        sd = ckpt[\"model_state_dict\"]\n",
    "    else:\n",
    "        # 有些直接保存了 state_dict 本身\n",
    "        sd = ckpt\n",
    "\n",
    "    # 列出参数名和形状\n",
    "    for k, v in sd.items():\n",
    "        try:\n",
    "            print(k, tuple(v.shape))\n",
    "        except Exception:\n",
    "            print(k, type(v))\n",
    "\n",
    "    # 统计参数量\n",
    "    total = sum(v.numel() for v in sd.values() if hasattr(v, \"numel\"))\n",
    "    print(\"total params:\", total)\n",
    "else:\n",
    "    print(\"checkpoint is not a dict, maybe a model object:\", ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "06cacd0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3633862/1611227685.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/both_bilstm.betterthanlast.b21_s15_epoch12.ckpt\", map_location=\"cpu\")\n"
     ]
    }
   ],
   "source": [
    "# 1) 载入 checkpoint（不要从不信任来源直接执行）\n",
    "ckpt = torch.load(\"/homeb/hpc/users/xyf/data/model/both_bilstm.betterthanlast.b21_s15_epoch12.ckpt\", map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ab7771b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'collections.OrderedDict'>\n",
      "keys: ['embed.weight', 'lstm_seq.weight_ih_l0', 'lstm_seq.weight_hh_l0', 'lstm_seq.bias_ih_l0', 'lstm_seq.bias_hh_l0', 'lstm_seq.weight_ih_l0_reverse', 'lstm_seq.weight_hh_l0_reverse', 'lstm_seq.bias_ih_l0_reverse', 'lstm_seq.bias_hh_l0_reverse', 'fc_seq.weight', 'fc_seq.bias', 'lstm_signal.weight_ih_l0', 'lstm_signal.weight_hh_l0', 'lstm_signal.bias_ih_l0', 'lstm_signal.bias_hh_l0', 'lstm_signal.weight_ih_l0_reverse', 'lstm_signal.weight_hh_l0_reverse', 'lstm_signal.bias_ih_l0_reverse', 'lstm_signal.bias_hh_l0_reverse', 'fc_signal.weight']\n",
      "embed.weight (16, 4)\n",
      "lstm_seq.weight_ih_l0 (512, 7)\n",
      "lstm_seq.weight_hh_l0 (512, 128)\n",
      "lstm_seq.bias_ih_l0 (512,)\n",
      "lstm_seq.bias_hh_l0 (512,)\n",
      "lstm_seq.weight_ih_l0_reverse (512, 7)\n",
      "lstm_seq.weight_hh_l0_reverse (512, 128)\n",
      "lstm_seq.bias_ih_l0_reverse (512,)\n",
      "lstm_seq.bias_hh_l0_reverse (512,)\n",
      "fc_seq.weight (128, 256)\n",
      "fc_seq.bias (128,)\n",
      "lstm_signal.weight_ih_l0 (512, 15)\n",
      "lstm_signal.weight_hh_l0 (512, 128)\n",
      "lstm_signal.bias_ih_l0 (512,)\n",
      "lstm_signal.bias_hh_l0 (512,)\n",
      "lstm_signal.weight_ih_l0_reverse (512, 15)\n",
      "lstm_signal.weight_hh_l0_reverse (512, 128)\n",
      "lstm_signal.bias_ih_l0_reverse (512,)\n",
      "lstm_signal.bias_hh_l0_reverse (512,)\n",
      "fc_signal.weight (128, 256)\n",
      "fc_signal.bias (128,)\n",
      "lstm_comb.weight_ih_l0 (1024, 256)\n",
      "lstm_comb.weight_hh_l0 (1024, 256)\n",
      "lstm_comb.bias_ih_l0 (1024,)\n",
      "lstm_comb.bias_hh_l0 (1024,)\n",
      "lstm_comb.weight_ih_l0_reverse (1024, 256)\n",
      "lstm_comb.weight_hh_l0_reverse (1024, 256)\n",
      "lstm_comb.bias_ih_l0_reverse (1024,)\n",
      "lstm_comb.bias_hh_l0_reverse (1024,)\n",
      "lstm_comb.weight_ih_l1 (1024, 512)\n",
      "lstm_comb.weight_hh_l1 (1024, 256)\n",
      "lstm_comb.bias_ih_l1 (1024,)\n",
      "lstm_comb.bias_hh_l1 (1024,)\n",
      "lstm_comb.weight_ih_l1_reverse (1024, 512)\n",
      "lstm_comb.weight_hh_l1_reverse (1024, 256)\n",
      "lstm_comb.bias_ih_l1_reverse (1024,)\n",
      "lstm_comb.bias_hh_l1_reverse (1024,)\n",
      "lstm_comb.weight_ih_l2 (1024, 512)\n",
      "lstm_comb.weight_hh_l2 (1024, 256)\n",
      "lstm_comb.bias_ih_l2 (1024,)\n",
      "lstm_comb.bias_hh_l2 (1024,)\n",
      "lstm_comb.weight_ih_l2_reverse (1024, 512)\n",
      "lstm_comb.weight_hh_l2_reverse (1024, 256)\n",
      "lstm_comb.bias_ih_l2_reverse (1024,)\n",
      "lstm_comb.bias_hh_l2_reverse (1024,)\n",
      "fc1.weight (256, 512)\n",
      "fc1.bias (256,)\n",
      "fc2.weight (2, 256)\n",
      "fc2.bias (2,)\n",
      "total params: 4693058\n"
     ]
    }
   ],
   "source": [
    "print(type(ckpt))\n",
    "\n",
    "# 2) 常见情况：Lightning 保存的 checkpoint 包含 'state_dict'\n",
    "if isinstance(ckpt, dict):\n",
    "    print(\"keys:\", list(ckpt.keys())[:20])\n",
    "    if \"state_dict\" in ckpt:\n",
    "        sd = ckpt[\"state_dict\"]\n",
    "    elif \"model_state_dict\" in ckpt:\n",
    "        sd = ckpt[\"model_state_dict\"]\n",
    "    else:\n",
    "        # 有些直接保存了 state_dict 本身\n",
    "        sd = ckpt\n",
    "\n",
    "    # 列出参数名和形状\n",
    "    for k, v in sd.items():\n",
    "        try:\n",
    "            print(k, tuple(v.shape))\n",
    "        except Exception:\n",
    "            print(k, type(v))\n",
    "\n",
    "    # 统计参数量\n",
    "    total = sum(v.numel() for v in sd.values() if hasattr(v, \"numel\"))\n",
    "    print(\"total params:\", total)\n",
    "else:\n",
    "    print(\"checkpoint is not a dict, maybe a model object:\", ckpt)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deepsignal3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
